import os,cv2, sys, shutil
from xml.dom.minidom import Document
convet2yoloformat = True
convert2vocformat = True

# 最小取20大小的脸，并且补齐
minsize2select = 20
usepadding = True

rootdir=r'/home/aistudio/work' #widerface数据集所在目录
datasetprefix=r'/home/aistudio/work/data' # 转换后数据地址
#
# rootdir=r'E:\dataset\wider' #widerface数据集所在目录
# datasetprefix=r'E:\dataset\wider\dev' # 转换后数据地址

imagesDir=datasetprefix + "/ImageSets"
if not os.path.exists(datasetprefix):
        os.mkdir(datasetprefix)

if not os.path.exists(imagesDir):
    os.mkdir(imagesDir)



def convertimgset(img_set="train"):
    imgdir = rootdir + "/WIDER_" + img_set + "/images"
    gtfilepath = rootdir + "/wider_face_split/wider_face_" + img_set + "_bbx_gt.txt"
    index = 0
    linesData=[]
    boxsData=[]

    with open(gtfilepath, 'r') as gtfile:
        while (True):  # and len(faces)<10
            filename = gtfile.readline()[:-1]
            if (filename == ""):
                break;
            if filename.find("0 0 0 0 0 0 0 0 0 0") >= 0:
                continue
            sys.stdout.write("\r" + str(index) + ":" + filename + "\t\t\t")
            sys.stdout.flush()
            imgpath = imgdir + "/" + filename
            img = cv2.imread(imgpath)
            if not img.data:
                break;
            imgheight = img.shape[0]
            imgwidth = img.shape[1]
            maxl = max(imgheight, imgwidth)
            paddingleft = (maxl - imgwidth) >> 1
            paddingright = (maxl - imgwidth) >> 1
            paddingbottom = (maxl - imgheight) >> 1
            paddingtop = (maxl - imgheight) >> 1
            saveimg = cv2.copyMakeBorder(img, paddingtop, paddingbottom, paddingleft, paddingright, cv2.BORDER_CONSTANT,
                                         value=0)
            numbbox = int(gtfile.readline())
            bboxes = []
            for i in range(numbbox):
                line = gtfile.readline()
                line = line.split()
                line = line[0:4]
                if (int(line[3]) <= 0 or int(line[2]) <= 0):
                    continue
                x = int(line[0]) + paddingleft
                y = int(line[1]) + paddingtop
                width = int(line[2])
                height = int(line[3])
                bbox = (x, y, width, height)
                x2 = x + width
                y2 = y + height
                # face=img[x:x2,y:y2]
                if width >= minsize2select and height >= minsize2select:
                    bboxes.append(bbox)

            filename = filename.replace("/", "_")
            if len(bboxes) == 0:
                print("warrning: no face")
                continue
            imageFullFile=imagesDir + "/" + filename
            cv2.imwrite(imageFullFile, saveimg)

            height = saveimg.shape[0]
            width = saveimg.shape[1]
            lineText="{} {} {} {}".format(index,imageFullFile,width,height)
            for i in range(len(bboxes)):
                bbox = bboxes[i]
                xcenter = (bbox[0] + bbox[2] * 0.5) / width
                ycenter = (bbox[1] + bbox[3] * 0.5) / height
                wr = bbox[2] * 1.0 / width
                hr = bbox[3] * 1.0 / height
                boxLine="0 " + str(xcenter) + " " + str(ycenter) + " " + str(wr) + " " + str(hr)
                lineText=lineText+" 0 {} {} {} {}".format(bbox[0],bbox[1],bbox[0]+bbox[2],bbox[1]+bbox[3])
                boxsData.append(boxLine)
            linesData.append(lineText)
            index = index + 1
    textFile = datasetprefix + "/" + img_set + ".txt"
    with open(textFile, 'w') as f:
        f.write("\n".join(linesData))

    labelFile = datasetprefix + "/label_" + img_set + ".txt"
    with open(labelFile, 'w') as f:
        f.write("\n".join(boxsData))


if __name__ == "__main__":
    img_sets = ["train","val"]
    for img_set in img_sets:
        convertimgset(img_set)